// Copyright 2024 Ant Group Co., Ltd.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#include "engine/operator/coalesce.h"

#include "gtest/gtest.h"

#include "engine/core/tensor_constructor.h"
#include "engine/operator/test_util.h"

namespace scql::engine::op {

struct CoalesceTestCase {
  std::vector<test::NamedTensor> exprs;
  test::NamedTensor expect_out;
};

class CoalesceTest
    : public testing::TestWithParam<
          std::tuple<test::SpuRuntimeTestCase, CoalesceTestCase>> {
 protected:
  static pb::ExecNode MakeExecNode(const CoalesceTestCase& tc);
  static void FeedInputs(ExecContext* ctx, const CoalesceTestCase& tc);
};

INSTANTIATE_TEST_SUITE_P(
    CoalesceBatchTest, CoalesceTest,
    testing::Combine(
        testing::Values(test::SpuRuntimeTestCase{spu::ProtocolKind::SEMI2K, 2}),
        testing::Values(
            // test private status
            CoalesceTestCase{
                .exprs = {test::NamedTensor(
                              "expr1", TensorFromJSON(arrow::int64(),
                                                      "[null, 2, 3, null]")),
                          test::NamedTensor(
                              "expr2", TensorFromJSON(arrow::int64(),
                                                      "[null, 12, 13, 14]"))},
                .expect_out = test::NamedTensor(
                    "out", TensorFromJSON(arrow::int64(), "[null, 2, 3, 14]"))},
            CoalesceTestCase{
                .exprs = {test::NamedTensor(
                    "expr", TensorFromJSON(arrow::float64(),
                                           "[null, -0.1, 1.1, null]"))},
                .expect_out = test::NamedTensor(
                    "out", TensorFromJSON(arrow::float64(),
                                          "[null, -0.1, 1.1, null]"))},
            CoalesceTestCase{
                .exprs =
                    {test::NamedTensor(
                         "expr1",
                         TensorFromJSON(arrow::large_utf8(),
                                        R"json(["B", null, null ,"B"])json")),
                     test::NamedTensor(
                         "expr2",
                         TensorFromJSON(arrow::large_utf8(),
                                        R"json(["CC", null, "CC" ,"C"])json"))},
                .expect_out = test::NamedTensor(
                    "out",
                    TensorFromJSON(arrow::large_utf8(),
                                   R"json(["B", null, "CC" ,"B"])json"))},
            CoalesceTestCase{
                .exprs = {test::NamedTensor(
                              "expr1", TensorFromJSON(arrow::float64(), "[]")),
                          test::NamedTensor(
                              "expr2", TensorFromJSON(arrow::float64(), "[]"))},
                .expect_out = test::NamedTensor(
                    "out", TensorFromJSON(arrow::float64(), "[]"))})),
    TestParamNameGenerator(CoalesceTest));

TEST_P(CoalesceTest, works) {
  // Given
  auto parm = GetParam();
  auto tc = std::get<1>(parm);
  auto node = MakeExecNode(tc);
  auto session = test::Make1PCSession();
  ExecContext ctx(node, session.get());

  FeedInputs(&ctx, tc);

  // When
  Coalesce op;
  ASSERT_NO_THROW(op.Run(&ctx));

  // Then
  // check alice output
  TensorPtr t = session->GetTensorTable()->GetTensor(tc.expect_out.name);
  ASSERT_TRUE(t);
  auto out_arr = t->ToArrowChunkedArray();

  auto expect_arr = tc.expect_out.tensor->ToArrowChunkedArray();

  // compare tensor content
  EXPECT_TRUE(out_arr->Equals(expect_arr))
      << "expect type = " << expect_arr->type()->ToString()
      << ", got type = " << out_arr->type()->ToString()
      << "\nexpect result = " << expect_arr->ToString()
      << "\nbut actual got result = " << out_arr->ToString();
}

/// ===================
/// CoalesceTest impl
/// ===================

pb::ExecNode CoalesceTest::MakeExecNode(const CoalesceTestCase& tc) {
  test::ExecNodeBuilder builder(Coalesce::kOpType);
  builder.SetNodeName("if-null-test");

  std::vector<pb::Tensor> inputs;
  for (auto expr : tc.exprs) {
    auto t = test::MakePrivateTensorReference(expr.name, expr.tensor->Type());
    inputs.push_back(t);
  }
  builder.AddInput(Coalesce::kExprs, inputs);

  {
    auto t = test::MakePrivateTensorReference(tc.expect_out.name,
                                              tc.expect_out.tensor->Type());
    builder.AddOutput(Coalesce::kOut, {t});
  }

  return builder.Build();
}

void CoalesceTest::FeedInputs(ExecContext* ctx, const CoalesceTestCase& tc) {
  test::FeedInputsAsPrivate(ctx, tc.exprs);
}

}  // namespace scql::engine::op